import xml.etree.ElementTree as ET
import numpy as np

class IQMsg(object):
    def __init__(self):
        self.width = None
        self.height = None

    def parse(self, xml_path):
        tree = ET.parse(xml_path)
        root = tree.getroot()

        ImageParam = root.find('ImageParam')
        self.height = int(ImageParam.find('Height').text)
        self.width = int(ImageParam.find('Width').text)

        PlayParam = root.find('PlayParam')
        self.Count = int(PlayParam.find('Count').text)


class VGridParam(object):
    def __init__(self):
        self.PContrast, self.DContrast, self.BContrast = 3.5, 3.5, 3.5
        self.brightness = 0.6
        self.SNRFlag, self.lpSubFlag, self.histFlag = True, False, False
        self.acqMode, self.acqChildmode, self.bodyPart = 1, 0, 1
        self.Fov, self.LAO_RAO, self.CRA_CAU, self.PlayMode = 20, 0, 0, 0


def getDiffBrightnessSequence(sequence, bright_list, vgrid_handle=None, tqdm_bar=None, vgParam: VGridParam()=None):
    _frame_count, _height, _width = sequence.shape
    if vgrid_handle is None or vgParam is None:
        print("warning: vgrid handle or vgParam not specified")
        return

    b_len = len(bright_list)
    result_list = []                # 最终的sequence list
    for b_id, BT in enumerate(bright_list):
        result_seq = []
        for frame_id, _img in enumerate(sequence):
            _img_out = np.zeros_like(_img)
            vgrid_handle.Process(
                _img, _img_out, vgParam.PContrast, vgParam.DContrast, vgParam.BContrast, 1.0, BT, 0,
                       vgParam.SNRFlag, vgParam.lpSubFlag, vgParam.histFlag, _width, _height, False,
                       vgParam.acqMode, vgParam.acqChildmode, vgParam.bodyPart,
                       20, 15, 30, -30, frame_id + 1, 0, False
                       )
            _img_out = np.clip(_img_out, 0, 65535)
            result_seq.append(_img_out)
            if tqdm_bar is not None:
                tqdm_bar.set_postfix_str(rf"B:{BT} {b_id + 1}/{b_len} {frame_id + 1}/{_frame_count}")
        result_list.append(result_seq)
    return np.array(result_list, dtype=np.uint16)







